import mplfinance as mpf
import pandas as pd
import numpy as np
import math


def writeFile(content):
  with open('test.txt', 'a') as file:
    file.write(content)

# 获取数据，处理数据
def read_data(filename):
  df = pd.read_csv(filename)
  df['Open'] = df['open']
  df['Close'] = df['close']
  df['High'] = df['high']
  df['Low'] = df['low']
  df['Volume'] = df['volume']
  df['Date'] = df['date']

  df['Date'] = pd.to_datetime(df['Date'])
  df.set_index(['Date'], inplace=True)
  return df

# 获取数据
# df = read_data('./files/history_A_stock_k_data.csv')
df = read_data('./files/茅台.csv')

df['MA5'] = df['Close'].rolling(window=5).mean()
df['MA10'] = df['Close'].rolling(window=10).mean()
df['MA20'] = df['Close'].rolling(window=20).mean()
df['MA60'] = df['Close'].rolling(window=60).mean()


df['Short_MA'] = df['Close'].rolling(window=12).mean()
df['Long_MA'] = df['Close'].rolling(window=26).mean()
df['DIF'] = df['Short_MA'] - df['Long_MA']        # 快线
df['DEA'] = df['DIF'].rolling(window=9).mean()    # 慢线
df['MACD'] = df['DIF'] - df['DEA']                # 柱子


df['macd-s'] = df['DIF']
df['macd-m'] = df['DEA']
df['macd-h'] = df['MACD']

obj = {
  'currentMoney': 100000,
  'currentNum': 0,
  'buyNumAll': 0,
  'sellNumAll': 0,
  'buyFeeAll': 0,
  'sellFeeAll': 0,
}
buyCommissionAll = 0
sellCommissionAll = 0

select = df[df['MA60'] > 0][['date', 'MA5','MA60']]
a = select[['date','MA5']]
b = select[['date','MA60']]
length = len(select)
list01 = []
list02 = []
for i in range(length):
  if(i != length - 1):
    if((a['MA5'].iloc[i] < b['MA60'].iloc[i]) and (a['MA5'].iloc[i + 1] > b['MA60'].iloc[i + 1])):
      list01.append(a['date'].iloc[i])

    if((a['MA5'].iloc[i] > b['MA60'].iloc[i]) and (a['MA5'].iloc[i + 1] < b['MA60'].iloc[i + 1])):
      list02.append(a['date'].iloc[i])
df['buy'] = df['date'].isin(list01)
df['sell'] = df['date'].isin(list02)

list03 = []
for i in range(length):
  if(i != length - 1):
    if((df['buy'].iloc[i]) or (df['sell'].iloc[i])):
      list03.append(df['date'].iloc[i])

finallyData = df[df['date'].isin(list03)][['date', 'close', 'MA5', 'MA60', 'buy', 'sell']]

finallyData.to_csv("./files/money.csv", index=False)

def buyAll(price, date):
  if(price * 100 > obj['currentMoney']):
    str = f'购买价格{price}，购买了0股份，还剩{obj['currentMoney']}元'
    print(str)
    writeFile(str + '\n')
  else:
    num = 0
    while True:
      currentNums = obj['currentNum']
      currentMoneys = obj['currentMoney']
      onePrice = price * 100
      currentNums = currentNums + (math.floor(currentMoneys / onePrice) - num) * 100
      buyMoney = currentNums * price
      buyCommission = buyMoney * 0.00025 + buyMoney * 0.00002
      if((buyMoney + buyCommission) > currentMoneys):
        num = num + 1
      else:
        obj['currentMoney'] = obj['currentMoney'] - buyMoney - buyCommission
        obj['currentNum'] = currentNums

        obj['buyNumAll'] = obj['buyNumAll'] + 1
        obj['buyFeeAll'] = obj['buyFeeAll'] + buyCommission

        str = f'购买日期{date}，购买价格{price}，购买了{obj['currentNum']}股份，还剩{obj['currentMoney']}元，手续费{buyCommission}'
        print(str)
        writeFile(str + '\n')
        break

def sellAll(price, date):
  nums = obj['currentNum']

  if(nums == 0):
    str = f'卖出价格{price}，卖了0股份，还剩{obj['currentMoney']}元'
    str02 = '--------------'
    print(str)
    writeFile(str + '\n')
    print(str02)
    writeFile(str02 + '\n')
  else:
    obj['currentNum'] = 0
    allMoney = nums * price
    buyCommission = allMoney * 0.00025 + allMoney * 0.00002 + allMoney * 0.001
    obj['currentMoney'] = obj['currentMoney'] + allMoney - buyCommission

    obj['sellNumAll'] = obj['sellNumAll'] + 1
    obj['sellFeeAll'] = obj['sellFeeAll'] + buyCommission
    
    str = f'卖出日期{date}，卖出价格{price}，卖出了{nums}股份，还剩{obj['currentMoney']}元，手续费{buyCommission}'
    str02 = '--------------'
    print(str)
    writeFile(str + '\n')
    print(str02)
    writeFile(str02 + '\n')


lengthB = len(finallyData)
for i in range(lengthB):
  if(i != lengthB - 1):
    if(finallyData['buy'].iloc[i]):
      buyAll(finallyData['close'].iloc[i], finallyData['date'].iloc[i])
    else:
      sellAll(finallyData['close'].iloc[i], finallyData['date'].iloc[i])

print('------------')
print(f'最终总余额：{obj['currentMoney']}')
print(f'最终总剩余股份：{obj['currentNum']}')
print(f'最终购买次数：{obj['buyNumAll']}')
print(f'最终总购买手续费：{obj['buyFeeAll']}')
print(f'最终卖出次数：{obj['sellNumAll']}')
print(f'最终总卖出手续费：{obj['sellFeeAll']}')

str = f'''
------------
最终总余额：{obj['currentMoney']}
最终总剩余股份：{obj['currentNum']}
最终购买次数：{obj['buyNumAll']}
最终总购买手续费：{obj['buyFeeAll']}
最终卖出次数：{obj['sellNumAll']}
最终总卖出手续费：{obj['sellFeeAll']}
'''
writeFile(str)





